#include "compi/compi.h"

#include <gtest/gtest.h>

#include <chrono>
#include <string>
#include <vector>

using namespace compi;
using namespace std::chrono;

// 全局测试环境指针，由自定义 main 显式创建与销毁，确保每个进程都调用 MPI_Init/Finalize。
static Environment * g_test_env = nullptr;

class SyncStrategiesTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        env_ = g_test_env; // 借用全局环境，统一 MPI 生命周期
        // 使用 Context::for_comm 获取非拥有的上下文实例
        comm_ = &Context::for_comm(MPI_COMM_WORLD, *env_);
    }

    void TearDown() override
    {
        // 非拥有指针清空，保留全局 Environment 直至进程退出
        comm_ = nullptr;
        env_ = nullptr;
    }

    Environment * env_{};
    Context * comm_{};
};

// 自定义 main：显式管理 MPI 生命周期，避免 mpiexec 认为未正确 finalize。
int
main(int argc, char ** argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    g_test_env = new Environment();
    const int code = RUN_ALL_TESTS();
    delete g_test_env;
    g_test_env = nullptr;
    return code;
}

// Ring 交换：Default 发送模式，验证无死锁且数据正确
TEST_F(SyncStrategiesTest, RingDefaultSend)
{
    int rank = comm_->rank();
    int size = comm_->size();
    int dest = (rank + 1) % size;
    int src = (rank - 1 + size) % size;

    auto & alloc = comm_->tag_allocator();
    auto ns = alloc.create_namespace("sync_ring_default");

    int recv_val = -1;
    MPI_Request rreq = MPI_REQUEST_NULL;
    MPI_Request sreq = MPI_REQUEST_NULL;
    MPI_Status status_recv{};

    auto t0 = steady_clock::now();

    // 先发布接收：使用 ANY_TAG，完成后校验命名空间归属
    int rc = MPI_Irecv(&recv_val, 1, MPI_INT, src, MPI_ANY_TAG, comm_->comm(), &rreq);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Irecv (RingDefaultSend)", comm_->comm()));

    // 发送到下一个进程（Default 模式，使用命名空间内标签）
    MessageTag send_tag = ns.request();
    int send_val = rank * 100 + dest;
    rc = MPI_Isend(&send_val, 1, MPI_INT, dest, static_cast<int>(send_tag), comm_->comm(), &sreq);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Isend (RingDefaultSend)", comm_->comm()));

    rc = MPI_Wait(&rreq, &status_recv);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Wait (RingDefaultSend rreq)", comm_->comm()));
    rc = MPI_Wait(&sreq, MPI_STATUS_IGNORE);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Wait (RingDefaultSend sreq)", comm_->comm()));

    // 释放发送标签资源
    ns.release(send_tag);

    auto t1 = steady_clock::now();

    // 校验接收到的数据与标签归属
    EXPECT_TRUE(ns.owns(status_recv.MPI_TAG));
    EXPECT_EQ(recv_val, src * 100 + rank);
    // 粗略时间上限，确保不会无限阻塞（环境差异下设定宽松阈值）
    EXPECT_LT(duration_cast<seconds>(t1 - t0).count(), 10);
}

// Ring 交换：Synchronous 发送模式，验证无死锁且数据正确
TEST_F(SyncStrategiesTest, RingSynchronousSend)
{
    int rank = comm_->rank();
    int size = comm_->size();
    int dest = (rank + 1) % size;
    int src = (rank - 1 + size) % size;

    auto & alloc = comm_->tag_allocator();
    auto ns = alloc.create_namespace("sync_ring_ssend");

    int recv_val = -1;
    MPI_Request rreq = MPI_REQUEST_NULL;
    MPI_Request sreq = MPI_REQUEST_NULL;
    MPI_Status status_recv{};

    auto t0 = steady_clock::now();

    // 关键：先发布接收，保证 Ssend 可匹配
    int rc = MPI_Irecv(&recv_val, 1, MPI_INT, src, MPI_ANY_TAG, comm_->comm(), &rreq);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Irecv (RingSynchronousSend)", comm_->comm()));

    // 使用同步发送（MPI_Issend），标签来自命名空间
    MessageTag send_tag = ns.request();
    int send_val = rank * 100 + dest;
    rc = MPI_Issend(&send_val, 1, MPI_INT, dest, static_cast<int>(send_tag), comm_->comm(), &sreq);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Issend (RingSynchronousSend)", comm_->comm()));

    rc = MPI_Wait(&rreq, &status_recv);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Wait (RingSynchronousSend rreq)", comm_->comm()));
    rc = MPI_Wait(&sreq, MPI_STATUS_IGNORE);
    ASSERT_TRUE(compi::call_mpi(rc, "MPI_Wait (RingSynchronousSend sreq)", comm_->comm()));

    // 释放发送标签
    ns.release(send_tag);

    auto t1 = steady_clock::now();

    EXPECT_TRUE(ns.owns(status_recv.MPI_TAG));
    EXPECT_EQ(recv_val, src * 100 + rank);
    EXPECT_LT(duration_cast<seconds>(t1 - t0).count(), 10);
}

// alltoall（固定类型）：Default 策略
TEST_F(SyncStrategiesTest, AlltoallFixedDefault)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<int> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = rank * 100 + j;

    auto t0 = steady_clock::now();
    auto recv_values =
        collectives::alltoall<int>(*comm_, send_values, SyncType::Default, SendMode::Default);
    auto t1 = steady_clock::now();

    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
        EXPECT_EQ(recv_values[i], i * 100 + rank);
    EXPECT_LT(duration_cast<seconds>(t1 - t0).count(), 10);
}

// alltoall（固定类型）：NBX + 同步发送
TEST_F(SyncStrategiesTest, AlltoallFixedNBXSync)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<int> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = rank * 100 + j;

    auto t0 = steady_clock::now();
    auto recv_values =
        collectives::alltoall<int>(*comm_, send_values, SyncType::NBX, SendMode::Synchronous);
    auto t1 = steady_clock::now();

    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
        EXPECT_EQ(recv_values[i], i * 100 + rank);
    EXPECT_LT(duration_cast<seconds>(t1 - t0).count(), 10);
}

// alltoall（变长字符串）：Default 策略
TEST_F(SyncStrategiesTest, AlltoallStringDefault)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<std::string> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = std::string("from_") + std::to_string(rank) + "_to_" + std::to_string(j);

    auto t0 = steady_clock::now();
    auto recv_values = collectives::alltoall<std::string>(
        *comm_, send_values, SyncType::Default, SendMode::Default);
    auto t1 = steady_clock::now();

    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
        EXPECT_EQ(recv_values[i],
                  std::string("from_") + std::to_string(i) + "_to_" + std::to_string(rank));
    EXPECT_LT(duration_cast<seconds>(t1 - t0).count(), 10);
}

// alltoall（变长字符串）：NBX + 同步发送
TEST_F(SyncStrategiesTest, AlltoallStringNBXSync)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<std::string> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = std::string("from_") + std::to_string(rank) + "_to_" + std::to_string(j);

    auto t0 = steady_clock::now();
    auto recv_values = collectives::alltoall<std::string>(
        *comm_, send_values, SyncType::NBX, SendMode::Synchronous);
    auto t1 = steady_clock::now();

    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
        EXPECT_EQ(recv_values[i],
                  std::string("from_") + std::to_string(i) + "_to_" + std::to_string(rank));
    EXPECT_LT(duration_cast<seconds>(t1 - t0).count(), 10);
}